#!/usr/bin/env bash

# Summarize TOP USER accounting information from Slurm sacct accounting records
# Author:	Ole.H.Nielsen@fysik.dtu.dk
# Homepage: https://github.com/OleHolmNielsen/Slurm_tools/

#####################################################################################
#
# Command usage:
#
function usage()
{
	cat <<EOF
Usage: slurmacct [-s Start_time -e End_time | -c | -w | -m monthyear] [-p partition(s)] [-u username] [-g groupname] [-G] [-W workdir] [-r report-prefix] [-n] [-h]
where:
	-s Start_time [last month]: Starting time of accounting period.
	-e End_time [last month]: End time of accounting period.
	-c: Current month
	-w: Last week
	-m monthyear: Select month and year (like "november2019")
	-p partition(s): Select only Slurm partion <partition>[,partition2,...]
	-u username: Print only user <username> 
	-g groupname: Print only users in UNIX group <groupname>
	-G: Print only groupwise summed accounting data
	-W directory: Print only jobs with this string in the job WorkDir
	-r: Report name prefix
	-n: No header information is printed
	-h: Print this help information

The Start_time and End_time values specify the date/time interval of
job completion/termination (see "man sacct").

Hint: Specify Start/End time as MMDD (Month and Date)
EOF
}

#####################################################################################

# Report file prefix
REPORT_PREFIX=/tmp/Slurm_report_acct_
export partition=""
export username=""
export groupname=""
export workdir=""
export ALLGROUPS=0
export SORTCOLUMN=5
export prefix=/usr/bin
export printheader=1
export month="last"

# Process options
while getopts "p:u:g:s:e:r:m:W:cwGhn" options; do
	case $options in
		p )     export partition=$OPTARG
			echo Print only accounting in Slurm partition $OPTARG
			;;
		u )     export username=$OPTARG
			export ALLGROUPS=0
			echo Print only user $OPTARG
			;;
		g )     export groupname="$OPTARG"
			export ALLGROUPS=0
	    		export SORTCOLUMN=5
			echo Print only users in UNIX group $OPTARG
			;;
		G )     export ALLGROUPS=1
			export username=""
			export groupname=""
	    		export SORTCOLUMN=4
			echo Print only groupwise summed accounting data
			;;
		s )     export start_time=$OPTARG
			echo Start date $OPTARG
			REPORT_NAME=${start_time}_${end_time}
			;;
		e )     export end_time=$OPTARG
			echo End date $OPTARG
			REPORT_NAME=${start_time}_${end_time}
			;;
		m )     echo Select month $OPTARG 
			start_time=`date -d "1$OPTARG" +%m01%y`
			end_time=`date -d "1$OPTARG + 1 month" +%m01%y`
			MONTH=`date -d "1$OPTARG" +%B`
			YEAR=`date -d "1$OPTARG" +%Y`
			REPORT_NAME=${MONTH}_${YEAR}
			;;
		c )     export start_time=`date +%m01%y`
			export end_time=`date +%m%d%y`
			echo Select current month from $start_time until $end_time
			REPORT_NAME=Current_month
			;;
		w )     export start_time=`date -d 'last week' +%m%d%y`
			export end_time=`date +%m%d%y`
			echo Select last week from $start_time until $end_time
			REPORT_NAME=Last_week
			;;
		W )     export workdir=$OPTARG
			echo Print only accounting data for jobs with working directory $workdir
			;;
		r )     export REPORT_PREFIX="$OPTARG"
			echo Copy report to $OPTARG
			;;
		n )     export printheader=0
			;;
		h|? ) usage
			exit 1;;
		* ) usage
			exit 1;;
	esac
done
shift $((OPTIND-1))

# Check Slurm commands path
if test -x /usr/local/bin/sacct
then
	export prefix=/usr/local/bin
elif ! test -x $prefix/sacct
then
	echo "ERROR: Slurm commands not found in '$prefix/' (Set 'prefix=')"
	exit 1
fi

#
# Default period: last month 
#
# Test if either start_time or end_time are empty strings
if test -z "$start_time" -o -z "$end_time"
then
	MONTH=`date -d "last month" +%B`
	YEAR=`date -d "last month" +%Y`
	REPORT_NAME=${MONTH}_${YEAR}
	start_time=`date -d "last month" +%m01%y`
	end_time=`date -d "last month + 1 month" +%m01%y`
fi

# Check partition names
if test -n "$partition"
then
	for p in `echo $partition | sed 's/,/ /g'`
	do
		# echo "Check partition $p"
	if test -z "`sinfo -h -p $p -O PartitionName`" 
	then
			echo "ERROR: Invalid partition name $p"
			# echo "Valid partition names are:"
			# sinfo -O "PartitionName"
			exit -1
		fi
	done
fi

# Test username
if test -n "$username" 
then
	if test -z "`$prefix/sacctmgr -p -n show assoc where users=$username`"
	then
		echo Error selecting Slurm username $username 
		exit -1
	fi
fi
# Test groupname
if test -n "$groupname"
then
	getent group $groupname > /dev/null
	if test "$?" -ne 0
	then
		echo "Error selecting UNIX groupname $groupname (it does not exist)"
		exit -1
	fi
fi

#####################################################################################
#
# Print a heading, and make selection for sacct report
#
# First report overall statistics including nicely formatted start/end date:
# Change the date/time format in report header for readibility (formats in "man strftime")

REPORT=${REPORT_PREFIX}${REPORT_NAME}

export SLURM_TIME_FORMAT="%d-%b-%Y_%R"

# Selections for sacct:
selection=""

# Print a report header

if test $printheader -gt 0
then
	# Print a sreport header (overwrite any existing file)
	$prefix/sreport cluster utilization start=$start_time end=$end_time -t percent > $REPORT
fi

# Request of a specific partition
if test -n "$partition"
then
	echo >> $REPORT
	echo Partition selected: $partition >> $REPORT
	selection="--partition $partition $selection"
fi

if test $ALLGROUPS -eq 0
then
	# User statistics
	echo >> $REPORT
	if test -n "$workdir"
	then
		echo "Print only accounting data for jobs with working directory $workdir" >> $REPORT
		echo >> $REPORT
	fi
	echo Usage sorted by top users: >> $REPORT
	# echo "Jobs completed/terminated between date/time $start_time and $end_time"
	if test -n "$username"
	then
		echo "User name selected: $username" >> $REPORT
		selection="--user $username $selection"
	else
		# Select all users
		selection="-a $selection"
	fi
	if test -n "$groupname"
	then
		echo "Group name selected: $groupname" >> $REPORT
		selection="--group $groupname $selection"
	fi
	echo "                             Wallclock          Average Average" >> $REPORT
	echo "Username    Group    #jobs       hours  Percent  #cpus  q-hours  Full name" >> $REPORT
	echo "--------    -----  ------- -----------  ------- ------- -------  ---------" >> $REPORT
else
	# Group statistics
	echo >> $REPORT
	if test -n "$workdir"
	then
		echo "Print only accounting data for jobs with working directory $workdir" >> $REPORT
		echo >> $REPORT
	fi
	echo Usage sorted by top groups: >> $REPORT
	# echo "Jobs completed/terminated between date/time $start_time and $end_time"
	echo "                    Wallclock          Average Average" >> $REPORT
	echo "   Group    #jobs       hours  Percent  #cpus  q-hours" >> $REPORT
	echo "   -----  ------- -----------  ------- ------- -------" >> $REPORT
fi

#####################################################################################
#
# Get and process Slurm accounting records

# Get length of strings for sacct formatting
if test -n "$usernamelength"
then
	ulen="%${#username}"
fi
if test -n "$groupnamelength"
then
	glen="%${#groupname}"
fi

# Report time in seconds:
export SLURM_TIME_FORMAT="%s"
# Request job data
export FORMAT="JobID,User${ulen},Group${glen},Partition,AllocNodes,AllocCPUS,Submit,Eligible,Start,End,CPUTimeRAW,State,WorkDir"
# Request job states: CAncelled, ReQueued, CompleteD, Failed, TimeOut, PReempted, Out_Of_Memory
export STATE="ca,rq,cd,f,to,pr,oom"

# Get Slurm individual job accounting records using the "sacct" command
# The "-a" flag was removed 04-Jan-2023
$prefix/sacct $selection -np -X -S $start_time -E $end_time -o $FORMAT -s $STATE | awk -F"|" '
BEGIN {
	userselect=ENVIRON["username"]
	groupselect=ENVIRON["groupname"]
	ALLGROUPS=ENVIRON["ALLGROUPS"]
	workdir=ENVIRON["workdir"]
	totaljobs=0
	# First get the list of user full names from /etc/passwd lines
	COMMAND="getent passwd"
	while (COMMAND | getline ) {
		split($0,b,":")		# Split password line into fields
		fullname[b[1]] = b[5]	# Full name b[5] of this username (b[1])
		# print b[1], fullname[b[1]]
	}
	close(COMMAND)
}
{
	# Parse input data
	JobID	= $1		# JobID
	cput	= $11		# CPU time in seconds
	if (cput <= 0) next	# Skip jobs with zero cputime
	user	= $2		# User name
	group	= $3		# Group name
	part	= $4		# Slurm partition name for this job
	nodect	= $5		# Number of nodes used
	total_ncpus = $6	# Total number of CPUs used (>=nodect)
	submit	= $7		# submit time
	eligible= $8		# eligible time
	if (eligible == "Unknown") eligible = submit
	start	= $9		# Job start time in epoch seconds
	end	= $10		# Job end time in epoch seconds
	if (start == "Unknown" || start == "None") start = end
	state	= $12		# Job state
	wall	= end - start
	wait	= start - eligible
	if (wait < 0) wait = 0 # Should not happen
	jobworkdir	= $13	# Job working directory

	# For accounting by number of CPU cores in stead of number of nodes,
	# uncomment the following line:
	nodect = total_ncpus

	# TOTAL accounting
	totaljobs++
	totalwait += wait
	cpunodesecs += nodect*cput
	wallnodesecs += nodect*wall
	wallsecs += wall

	# Check if this job matches selection criteria
	if (userselect != "" && user != userselect) next
	if (groupselect != "" && group != groupselect) next
	if (workdir != "" && jobworkdir !~ workdir) next

	# User accounting (a user may belong to several groups)
	usernamegroup[user,group] = user
	fullnamegroup[user,group] = fullname[user]
	usergroup[user,group] = group
	jobs[user,group]++
	cpunodes[user,group] += nodect*cput
	wallnodes[user,group] += nodect*wall
	wallcpu[user,group] += wall
	if (nodect < minnodes[user,group]) minnodes[user,group] = nodect
	if (nodect > maxnodes[user,group]) maxnodes[user,group] = nodect
	waittime[user,group] += wait

	# Group accounting
	groupname[group]=group
	gr_jobs[group]++
	gr_cpunodes[group] += nodect*cput
	gr_wallnodes[group] += nodect*wall
	gr_wallcpu[group] += wall
	if (nodect < gr_minnodes[group]) gr_minnodes[group] = nodect
	if (nodect > gr_maxnodes[group]) gr_maxnodes[group] = nodect
	gr_waittime[group] += wait
} END {
	# Some average values
	if (totaljobs > 0)
		totalwaitaverage = totalwait/totaljobs
	else
		totalwaitaverage = 0
	if (wallsecs > 0)
		nodesaverage = wallnodesecs/wallsecs
	else
		nodesaverage = 0
	# Usage in hours
	cpunodehours = cpunodesecs / 3600
	wallnodehours = wallnodesecs / 3600
	wallhours = wallsecs / 3600
	if (totaljobs == 0 || cpunodehours == 0 || wallnodehours == 0 || wallhours == 0) {
		# print "ERROR: Zero CPU hours recorded in specified date interval"
		# print "totaljobs = ", totaljobs, "cpunodehours = ", cpunodehours, "wallnodehours = ", wallnodehours, "wallhours = ", wallhours
		# exit 1
	}

	if (ALLGROUPS == 0) {

		# Print format
		format = "%8.8s %8.8s %8d  %10.1f  %7.2f %7.2f %7.2f  %s\n"
		groupjobs = 0
		grouphours = 0
		for (ug in usernamegroup) {
			if (length(groupselect) > 0 && usergroup[ug] != groupselect) continue
			if (wallnodehours > 0 && jobs[ug] > 0 && wallcpu[ug] > 0)
				printf(format, usernamegroup[ug], usergroup[ug],
				jobs[ug], 
				wallnodes[ug]/3600, 
				100*wallnodes[ug]/(wallnodehours*3600),
				wallnodes[ug]/wallcpu[ug],
				waittime[ug]/jobs[ug]/3600,
				fullnamegroup[ug])
			groupjobs += jobs[ug]
			groupnodehours += wallnodes[ug]/3600
			grouphours += wallcpu[ug]/3600
			groupwait += waittime[ug]
			usercount++
		}
		# Print out total usage
		totalusers = "Number of users: " usercount
		printf(format, "TOTAL", "(All)", totaljobs, wallnodehours, 100,
			nodesaverage, totalwaitaverage/3600, totalusers)
		# Print out group usage
		if (length(groupselect) > 0 && groupjobs > 0 && wallnodehours > 0 && grouphours > 0 && groupjobs > 0)
			printf(format,
				"GROUP", groupselect, groupjobs, groupnodehours,
				100*groupnodehours/wallnodehours,
				groupnodehours/grouphours, groupwait/groupjobs/3600, "")
	} else {

		# Per-group accounting

		# Print format
		format = "%8s %8d  %10.1f  %7.2f %7.2f %7.2f\n"

		# Sort arrays by element values:
		# https://www.gnu.org/software/gawk/manual/html_node/Controlling-Scanning.html
		PROCINFO["sorted_in"] = "@ind_num_desc"

		for (group in groupname) {
			if (gr_jobs[group] > 0 && gr_wallcpu[group] > 0) printf(format,
				groupname[group],
				gr_jobs[group], 
				gr_wallnodes[group]/3600,
				100*gr_wallnodes[group]/(3600*wallnodehours),
				gr_wallnodes[group]/gr_wallcpu[group],
				gr_waittime[group]/gr_jobs[group]/3600)
		}
		printf(format, "TOTAL", totaljobs, wallnodehours, 100,
			nodesaverage, totalwaitaverage/3600)
	}
		
} ' | env LC_ALL=C sort -r -n -k $SORTCOLUMN -k 3 -k 1d >> $REPORT

# The sort command sorts number of running procs in descending order
# on keys $SORTCOLUMN and 3, and alphabetical sort on key 1
# The LC_ALL=C ensures that Upper case is sorted before lower case.

echo Report generated to file $REPORT

exit 0
